# -*- coding: utf-8 -*-
import numpy as np
import os, cv2, time
# from config_loader import *
from tools.utils import *
import imgaug.augmenters as iaa
from skimage.morphology import skeletonize

def aug_ink_skin_rgb(input_image, defect_num=(2, 3), thickness=(3, 7), gray_threshold=40):
    """
    说明：墨皮仿真
    """
    if input_image.ndim == 2:
        src_image = cv2.cvtColor(input_image, cv2.COLOR_GRAY2BGR)
    else:
        src_image = input_image.copy()

    # 获取图像尺寸
    height, width = src_image.shape[:2]

    # 创建一个与输入图像相同大小的透明图层（3通道：BGR）
    overlay = np.zeros_like(src_image)

    # 创建用于返回的掩膜图层
    defect_mask = np.zeros((height, width), dtype=np.uint8)

    # 随机生成一些半透明的环形椭圆或圆形作为墨皮缺陷
    min_defect_num, max_defect_num = defect_num
    random_num_defects = np.random.randint(min_defect_num, max_defect_num)  # 随机生成缺陷数量

    for _ in range(random_num_defects):
        # 随机选择缺陷的位置、大小和颜色
        center_x, center_y = np.random.randint(0, width), np.random.randint(0, height)
        outer_axis_major = np.random.randint(3, 12)  # 外层椭圆主轴长度
        outer_axis_minor = np.random.randint(4, 12)  # 外层椭圆副轴长度
        min_thickness, max_thickness = thickness
        rand_thickness = np.random.randint(min_thickness, max_thickness)

        # 内层椭圆的轴长比外层小一个固定的厚度
        inner_axis_major = max(1, outer_axis_major - rand_thickness)
        inner_axis_minor = max(1, outer_axis_minor - rand_thickness)

        # 创建临时掩膜用于判断该区域是否满足条件
        temp_mask = np.zeros((height, width), dtype=np.uint8)

        # 在临时掩膜上绘制外层椭圆（实心）
        cv2.ellipse(temp_mask, (center_x, center_y), (outer_axis_major, outer_axis_minor),
                    angle=0, startAngle=0, endAngle=360, color=255, thickness=-1)
        # 绘制内层椭圆（实心），覆盖掉中间部分以形成环形
        cv2.ellipse(temp_mask, (center_x, center_y), (inner_axis_major, inner_axis_minor),
                    angle=0, startAngle=0, endAngle=360, color=0, thickness=-1)

        # 获取环形区域的坐标
        ys, xs = np.where(temp_mask == 255)

        # 如果没有有效像素，跳过这个缺陷
        if len(xs) == 0 or len(ys) == 0:
            continue

        # 检查坐标是否在图像范围内
        valid_xs = xs[(xs >= 0) & (xs < width)]
        valid_ys = ys[(ys >= 0) & (ys < height)]

        if len(valid_xs) == 0 or len(valid_ys) == 0:
            continue
        # 随机选择半透明的颜色 (B, G, R)
        color = (
            np.random.randint(100, 120),  # 蓝色通道
            np.random.randint(100, 121),  # 绿色通道
            np.random.randint(102, 110),  # 红色通道
        )
        # 计算原图在该区域的平均灰度值
        mean_image_gray = np.median(src_image[valid_ys, valid_xs, :])
        # 只有当原图该区域的灰度值大于阈值时，才添加该缺陷
        if mean_image_gray < 220:
            # 在overlay上绘制外层椭圆（实心）
            cv2.ellipse(overlay, (center_x, center_y), (outer_axis_major, outer_axis_minor),
                        angle=0, startAngle=0, endAngle=360, color=color, thickness=-1)
            # 绘制内层椭圆（实心），覆盖掉中间部分以形成环形
            cv2.ellipse(overlay, (center_x, center_y), (inner_axis_major, inner_axis_minor),
                        angle=0, startAngle=0, endAngle=360, color=(0, 0, 0), thickness=-1)

            # 在defect_mask上绘制白色的环形区域
            cv2.ellipse(defect_mask, (center_x, center_y), (outer_axis_major, outer_axis_minor),
                        angle=0, startAngle=0, endAngle=360, color=255, thickness=-1)
            cv2.ellipse(defect_mask, (center_x, center_y), (inner_axis_major, inner_axis_minor),
                        angle=0, startAngle=0, endAngle=360, color=0, thickness=-1)
    if np.count_nonzero(defect_mask) == 0:
        return None, None, False

    # 将叠加层与原图混合
    alpha = np.random.uniform(0.35, 0.814)
    result_image = cv2.addWeighted(overlay, alpha, src_image, 1.0, 0)

    # 外边缘羽化修复
    canny_mask = cv2.Canny(overlay, 100, 200)
    result_image = feather_region_with_mask(result_image, canny_mask, feather_radius=2)

    # 确保输入和输出结果一致
    if input_image.ndim == 2:
        result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2GRAY)
    return result_image, defect_mask, True


def aug_text_blur_rgb(input_image,
                  text_mask,
                  defect_type_list=['circle', 'ellipse', 'polygon'],
                  defect_num=(3, 5),
                  pt_size=(1, 10)):
    """
    说明： 这个函数仅仅对文本区域进行丢失仿真，针对背景或其他区域都不生成任何缺陷
    参数：
    input_image: 输入的待仿真的检测图片
    atext_mask: 与输入对应的文本mask， 白色：文本行， 黑色：背景
    defect_type_list: ['line', 'circle', 'ellipse', 'polygon']
    pt_size: 指的是缺陷半径
    """
    # 如果是单通道图像，需要转换成RGB图像
    if input_image.ndim == 2:
        input_image_cpy = cv2.cvtColor(input_image, cv2.COLOR_GRAY2BGR)
    else:
        input_image_cpy = input_image.copy()
    if text_mask.ndim == 3:
        text_mask = cv2.cvtColor(text_mask, cv2.COLOR_BGR2GRAY)

    gray_image = cv2.cvtColor(input_image_cpy, cv2.COLOR_BGR2GRAY)
    image_h, image_w = input_image_cpy.shape[0:2]

    # 创建一个与图像大小相同的空白图像，用于存储墨点
    defect_mask = np.zeros_like(input_image, dtype=np.uint8)

    exact_text_mask = text_mask_extact(gray_image, text_mask)
    binary_img = copy.deepcopy(exact_text_mask)
    binary_bool_img = binary_img.astype(bool)
    skeleton_img = skeletonize(binary_bool_img)
    skeleton_img = ((skeleton_img > 0).astype(np.uint8))*255
    gen_seed_y, gen_seed_x = np.where(skeleton_img == 255)

    # 如果当前图片没有文本框，直接返回
    is_gen_defect = True  # 判断给定的输入图像是否适合输入图像标志位
    if len(gen_seed_x) == 0:
        is_gen_defect = False
        return None, None, is_gen_defect  # 如果没有文本框，就不存在生成缺陷的可能，直接返回

    # 获取生成缺陷的数量
    min_defect_num, max_defect_num = defect_num
    if min_defect_num == max_defect_num:
        num_defect = min_defect_num
    else:
        num_defect = random.randint(min_defect_num, max_defect_num)

    # 生成随机缺陷
    bright_image = aug_light(input_image, gamma=(0.1, 0.2))
    dark_image = input_image.copy()
    for _ in range(num_defect):
        # 随机获得生成缺陷的中心点坐标
        gen_seed_id = random.randint(1, len(gen_seed_x)-1)
        gen_center_x, gen_center_y = gen_seed_x[gen_seed_id], gen_seed_y[gen_seed_id]
        gen_center_gray_value = gray_image[gen_center_y, gen_center_x]

        # 文字经过精修，直接判断文字灰度
        if gen_center_gray_value > 128:
            texture_flag = 0  # 黑底白字
        else:
            texture_flag = 1  # 白底黑字

        # 随机确定脏点半径
        min_pt_size, max_pt_size = pt_size
        random_pt_size = random.randint(min_pt_size, max_pt_size)

        # 选择随机形状
        spot_type = random.choice(defect_type_list)
        temp_mask = np.zeros((image_h, image_w), np.uint8)
        if spot_type == 'circle':
            cv2.circle(temp_mask, (gen_center_x, gen_center_y), random_pt_size, 255, -1)

        elif spot_type == 'ellipse':
            axes = (random.randint(random_pt_size // 2, random_pt_size),
                    random.randint(random_pt_size // 2, random_pt_size))
            cv2.ellipse(temp_mask, (gen_center_x, gen_center_y), axes, 0, 0, 360, 255, -1)

        elif spot_type == 'polygon':
            polygon_points = generate_polygon((gen_center_x, gen_center_y), random_pt_size,
                                              irregularity=0.5, spikeyness=0.4, num_vertices=5)
            cv2.fillPoly(temp_mask, [np.int32(polygon_points)], 255)

        elif spot_type == 'line':
            pt1 = (random.randint(0, image_w - 1), random.randint(0, image_h - 1))
            pt2 = (random.randint(0, image_w - 1), random.randint(0, image_h - 1))
            thickness = random.randint(1, 2)  # 随机厚度
            cv2.line(temp_mask, pt1, pt2,  255, thickness)

        # 形状与文字mask做与运算，得到精准的文字区域（只想把文字的亮度变化，背景不变）
        text_binary = cv2.bitwise_and(exact_text_mask, temp_mask)
        white_y, white_x = np.where(text_binary == 255)
        if texture_flag == 0:  # 0：黑底白字
            white_number = np.where(gray_image[white_y, white_x] >= 128)
            text_ratio = len(white_number[0]) / (len(white_y) + 1)
        else:  # 1：白底黑字
            black_number = np.where(gray_image[white_y, white_x] < 128)
            text_ratio = len(black_number[0]) / (len(white_y) + 1)

        if text_ratio < 0.05:
            continue
        else:
            if texture_flag == 0:  # 0：黑底白字
                mean_val_ori = np.median(input_image[white_y, white_x, :])
                mean_val_dst = np.median(dark_image[white_y, white_x, :])
                if abs(mean_val_dst - mean_val_ori) > 30:
                    input_image_cpy[white_y, white_x, :] = dark_image[white_y, white_x, :].copy()
                    defect_mask[white_y, white_x, :] = (255, 255, 255)
            else:
                # 1：白底黑字
                mean_val_ori = np.median(input_image[white_y, white_x, :])
                mean_val_dst = np.median(bright_image[white_y, white_x, :])
                if abs(mean_val_dst - mean_val_ori) > 30:
                    input_image_cpy[white_y, white_x, :] = bright_image[white_y, white_x, :].copy()
                    defect_mask[white_y, white_x, :] = (255, 255, 255)

    if np.count_nonzero(defect_mask != 0) == 0:
        is_gen_defect = False
        return None, None, is_gen_defect  # 如果没有文本框，就不存在生成缺陷的可能，直接返回

    # 对添加的区域进行边缘提取，便于把生成的缺陷边缘羽化，让其显示的更加自然
    canny_mask = cv2.Canny(defect_mask, 13, 13)
    result_image = feather_region_with_mask(input_image_cpy, canny_mask, feather_radius=7)
    # result_image = input_image_cpy.copy()

    # 确保输入和输出结果一致
    if input_image.ndim == 2:
        result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2GRAY)
    mask = cv2.cvtColor(defect_mask, cv2.COLOR_BGR2GRAY)
    return result_image, mask, is_gen_defect


def aug_text_miss_rgb(input_image,
                  text_mask,
                  defect_type_list=['circle', 'ellipse', 'polygon'],
                  defect_num=(3, 5),
                  pt_size=(5, 10)):
    """
    说明： 这个函数仅仅对文本区域进行丢失仿真，针对背景或其他区域都不生成任何缺陷
    参数：
    input_image: 输入的待仿真的检测图片
    atext_mask: 与输入对应的文本mask， 白色：文本行， 黑色：背景
    defect_type_list: ['line', 'circle', 'ellipse', 'polygon']
    pt_size: 指的是缺陷半径
    """
    # 1. 保留原始彩色图像，不再转灰度
    if input_image.ndim == 3:
        input_image_cpy = input_image.copy()
    else:
        input_image_cpy = cv2.cvtColor(input_image, cv2.COLOR_GRAY2BGR)
    gray_image = cv2.cvtColor(input_image_cpy, cv2.COLOR_BGR2GRAY)
    # 2. 确保 text_mask 是单通道
    if text_mask.ndim == 3:
        text_mask = cv2.cvtColor(text_mask, cv2.COLOR_BGR2GRAY)
    image_h, image_w = input_image_cpy.shape[:2]

    # 获取白底黑字还是黑底白字的标志位（1：白底黑字。 0：黑底白字。）
    texture_flag = get_background_flag(gray_image)

    # 获取文字的基础颜色，便于用这个颜色生成缺陷
    rand_color_list, _ = get_background_mean_val(input_image, texture=texture_flag)
    r, g, b = rand_color_list[0:3]

    # 创建一个与图像大小相同的空白图像，用于存储墨点
    defect_mask = np.zeros_like(input_image, dtype=np.uint8)

    # 生成缺陷生成位置的种子点（是在背景上生成，还是在前景上生成）
    # - 获取文本mask区域的灰度纹理
    gray_image = cv2.bilateralFilter(gray_image, d=9, sigmaColor=25, sigmaSpace=25)
    gray_text_image = cv2.bitwise_and(gray_image, text_mask)

    if texture_flag == 0:  # 0：黑底白字
        _, binary_img = cv2.threshold(gray_text_image, 128, 255, cv2.THRESH_BINARY)
    else:  # 1：白底黑字
        gray_text_image[gray_text_image < 5] = 255
        _, binary_img = cv2.threshold(gray_text_image, 128, 255, cv2.THRESH_BINARY_INV)

    binary_bool_img = binary_img.astype(bool)
    skeleton_img = skeletonize(binary_bool_img)
    skeleton_img = ((skeleton_img > 0).astype(np.uint8))*255
    gen_seed_y, gen_seed_x = np.where(skeleton_img == 255)

    # 如果当前图片没有文本框，直接返回
    is_gen_defect = True  # 判断给定的输入图像是否适合输入图像标志位
    if len(gen_seed_x) == 0:
        is_gen_defect = False
        return None, None, is_gen_defect  # 如果没有文本框，就不存在生成缺陷的可能，直接返回

    # 获取生成缺陷的数量
    min_defect_num, max_defect_num = defect_num
    if min_defect_num == max_defect_num:
        num_defect = min_defect_num
    else:
        num_defect = random.randint(min_defect_num, max_defect_num)

    # 生成随机缺陷
    for _ in range(num_defect):
        # 随机获得生成缺陷的中心点坐标
        gen_seed_id = random.randint(1, len(gen_seed_x)-1)
        gen_center_x, gen_center_y = gen_seed_x[gen_seed_id], gen_seed_y[gen_seed_id]

        # 随机确定脏点半径
        min_pt_size, max_pt_size = pt_size
        random_pt_size = random.randint(min_pt_size, max_pt_size)

        # 生成随机颜色
        if texture_flag == 1:  # 白底黑字
            rand_gray = random.randint(1, 10)
            color = (max(0, min(255, r + rand_gray)),
                     max(0, min(255, g + rand_gray)),
                     max(0, min(255, b + rand_gray)))
        else:
            rand_gray = random.randint(1, 10)
            color = (max(0, min(255, r - rand_gray)),
                     max(0, min(255, g - rand_gray)),
                     max(0, min(255, b - rand_gray)))

        # 选择随机形状
        spot_type = random.choice(defect_type_list)
        temp_mask = np.zeros((image_h, image_w), np.uint8)
        if spot_type == 'circle':
            cv2.circle(temp_mask, (gen_center_x, gen_center_y), random_pt_size, 255, -1)

        elif spot_type == 'ellipse':
            axes = (random.randint(random_pt_size // 2, random_pt_size),
                    random.randint(random_pt_size // 2, random_pt_size))
            cv2.ellipse(temp_mask, (gen_center_x, gen_center_y), axes, 0, 0, 360, 255, -1)

        elif spot_type == 'polygon':
            polygon_points = generate_polygon((gen_center_x, gen_center_y), random_pt_size,
                                              irregularity=0.5, spikeyness=0.4, num_vertices=5)
            cv2.fillPoly(temp_mask, [np.int32(polygon_points)], 255)

        elif spot_type == 'line':
            pt1 = (random.randint(0, image_w - 1), random.randint(0, image_h - 1))
            pt2 = (random.randint(0, image_w - 1), random.randint(0, image_h - 1))
            thickness = random.randint(1, 2)  # 随机厚度
            cv2.line(temp_mask, pt1, pt2,  255, thickness)

        # 4-29新增： 针对仿真区域进行灰度判断， 如果新增的灰度与原始图像区域的灰度相差太小， 这个仿真就不能再添加了
        white_y, white_x = np.where(temp_mask == 255)
        if texture_flag == 0:  # 0：黑底白字
            white_number = np.where(gray_image[white_y, white_x] >= 128)
            text_ratio = len(white_number[0]) / (len(white_y) + 1)
        else:  # 1：白底黑字
            black_number = np.where(gray_image[white_y, white_x] < 128)
            text_ratio = len(black_number[0]) / (len(white_y) + 1)

        if text_ratio < 0.05:
            continue
        else:
            defect_mask[white_y, white_x, :] = (255, 255, 255)
            input_image_cpy[white_y, white_x, :] = color

    if np.count_nonzero(defect_mask != 0) == 0:
        is_gen_defect = False
        return None, None, is_gen_defect  # 如果没有文本框，就不存在生成缺陷的可能，直接返回

    # 对添加的区域进行边缘提取，便于把生成的缺陷边缘羽化，让其显示的更加自然
    canny_mask = cv2.Canny(defect_mask, 13, 13)
    result_image = feather_region_with_mask(input_image_cpy, canny_mask, feather_radius=5)

    # 确保输入和输出结果一致
    if input_image.ndim == 2:
        result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2GRAY)

    mask = cv2.cvtColor(defect_mask, cv2.COLOR_BGR2GRAY)
    return result_image, mask, is_gen_defect


def aug_dirty_pt_rgb(input_image,
                       text_mask,
                       is_gen_background=True,
                       defect_type_list=['circle', 'ellipse', 'polygon', 'line'],
                       defect_num=(5, 10),
                       pt_size=(4, 5),
                       line_size=(10, 100)):
    """
    在彩色印刷品文本区域周边生成带有真实颜色的脏点、划痕等缺陷。
    """
    # 1. 保留原始彩色图像，不再转灰度
    if input_image.ndim == 3:
        input_image_cpy = input_image.copy()
    else:
        input_image_cpy = cv2.cvtColor(input_image, cv2.COLOR_GRAY2BGR)

    # 2. 确保 text_mask 是单通道
    if text_mask.ndim == 3:
        text_mask = cv2.cvtColor(text_mask, cv2.COLOR_BGR2GRAY)
    image_h, image_w = input_image_cpy.shape[:2]

    # 3. 获取背景类型（白底黑字 or 黑底白字）
    #    用原始图像的均值判断（更准确）
    mean_color = np.mean(input_image_cpy, axis=2)  # 转为近似灰度
    texture_flag = 1 if np.mean(mean_color) > 128 else 0  # 1: 白底，0: 黑底

    # 4. 确定缺陷生成种子点
    if is_gen_background:
        # 背景区域：mask == 0
        gen_seed_y, gen_seed_x = np.where(text_mask == 0)
    else:
        # 前景文本边缘：提取文本边缘，并膨胀一点避免贴边
        gray_text_image = cv2.bitwise_and(mean_color.astype(np.uint8), text_mask)
        canny_mask = cv2.Canny(gray_text_image, 33, 33)
        kernel = np.ones((7, 7), np.uint8)
        canny_mask_dilate = cv2.dilate(canny_mask, kernel, iterations=2)
        gen_seed_y, gen_seed_x = np.where(canny_mask_dilate == 255)

    if len(gen_seed_x) == 0:
        return None, None, False

    # 5. 随机生成缺陷数量
    min_defect_num, max_defect_num = defect_num
    num_defect = random.randint(min_defect_num, max_defect_num)

    # 6. 初始化缺陷掩膜
    defect_mask = np.zeros((image_h, image_w), dtype=np.uint8)

    for _ in range(num_defect):
        # 随机获得生成缺陷的中心点坐标
        gen_seed_id = random.randint(1, len(gen_seed_x) - 1)
        gen_center_x, gen_center_y = gen_seed_x[gen_seed_id], gen_seed_y[gen_seed_id]
        # 随机确定脏点半径
        random_pt_size = random.randint(pt_size[0], pt_size[1])
        temp_mask = np.zeros((image_h, image_w), np.uint8)
        spot_type = random.choice(defect_type_list)

        # 选择随机形状
        if spot_type == 'circle':
            cv2.circle(temp_mask, (gen_center_x, gen_center_y), random_pt_size, 255, -1)
        elif spot_type == 'ellipse':
            axes = (random.randint(random_pt_size // 2, random_pt_size),
                    random.randint(random_pt_size // 2, random_pt_size))
            cv2.ellipse(temp_mask, (gen_center_x, gen_center_y), axes, 0, 0, 360, 255, -1)
        elif spot_type == 'polygon':
            polygon_points = generate_polygon((gen_center_x, gen_center_y), random_pt_size,
                                              irregularity=0.5, spikeyness=0.4, num_vertices=5)
            cv2.fillPoly(temp_mask, [np.int32(polygon_points)], 255)
        elif spot_type == 'line':
            min_pt_size, max_pt_size = line_size
            line_length = random.randint(min_pt_size, max_pt_size)

            # 随机获得生成缺陷的中心点坐标
            gen_seed_id = random.randint(0, len(gen_seed_x) - 1)  # 修改为从0开始以包含数组的第一个元素
            gen_center_x, gen_center_y = gen_seed_x[gen_seed_id], gen_seed_y[gen_seed_id]

            # 计算方向角并转换为弧度
            angle = np.random.uniform(0, 2 * np.pi)

            # 根据角度和指定长度计算终点
            end_x = int(gen_center_x + line_length * np.cos(angle))
            end_y = int(gen_center_y + line_length * np.sin(angle))

            # 调整终点以确保其在图像范围内
            end_x = max(0, min(image_w - 1, end_x))
            end_y = max(0, min(image_h - 1, end_y))

            thickness = random.randint(1, 2)  # 随机厚度
            temp_mask = np.zeros((image_h, image_w), np.uint8)
            cv2.line(temp_mask, (gen_center_x, gen_center_y), (end_x, end_y), 255, thickness, cv2.LINE_4)

        # 获取缺陷区域像素
        white_y, white_x = np.where(temp_mask > 0)
        if len(white_x) == 0:
            continue

        # 7. 获取该区域原始平均颜色（BGR）
        roi_colors = input_image_cpy[white_y, white_x]  # shape: (N, 3)
        mean_bgr = np.median(roi_colors, axis=0).astype(int)  # [B, G, R]

        # 8. 生成“脏色”——模拟串色、氧化、油墨污染
        new_color = np.zeros(3, dtype=int)
        for ch in range(3):  # 对 B, G, R 通道分别处理
            # 从原始颜色 ±30 以外选一个值，制造“异常色”
            new_color[ch] = random_excluding(
                center=mean_bgr[ch],
                exclude_range=30,
                min_val=0,
                max_val=255
            )

        # 9. 可选：让颜色更“脏”一点（比如偏棕、偏灰）
        # 例如：模拟氧化（偏黄/棕）: R↑, G↑, B↓
        if random.random() < 0.8:  # 60% 概率添加脏色调
            new_color[0] = max(0, new_color[0] - random.randint(20, 50))  # 减少蓝色
            new_color[1] = min(255, new_color[1] + random.randint(10, 30))  # 增加绿色
            # 红色可保持或增加
            new_color = np.clip(new_color, 0, 255).astype(int)

        # 10. 写入新颜色
        input_image_cpy[white_y, white_x] = new_color
        defect_mask[white_y, white_x] = 255

    # 11. 边缘羽化，让缺陷更自然
    canny_mask = cv2.Canny(defect_mask, 13, 13)
    result_image = feather_region_with_mask(input_image_cpy, canny_mask, feather_radius=7)

    # 12. 返回结果
    if np.count_nonzero(defect_mask) == 0:
        return None, None, False
    else:
        return result_image, defect_mask, True


def aug_color_miss_plus_rgb(input_image,
                        defect_mask_root,
                        defect_num=(5, 6),
                        gray_range=(85, 90),  #gray_range=(35, 60),
                        scale_range=(0.2, 1.5),
                        rotate=True, distort=True):
    # 获取所有缺陷掩膜图片路径
    defect_mask_path_list = get_file_path_list(defect_mask_root, FileType.IMAGE_TYPE)

    # 判断输入是否为彩色图
    if input_image.ndim == 3:
        input_image_cpy = input_image.copy()
    else:
        input_image_cpy = cv2.cvtColor(input_image, cv2.COLOR_GRAY2BGR)  # 转换为三通道以支持彩色输出
    image_h, image_w = input_image_cpy.shape[:2]

    # 获取生成缺陷的数量
    min_defect_num, max_defect_num = defect_num
    num_defect = random.randint(min_defect_num, max_defect_num)
    defect_mask = np.zeros((image_h, image_w), np.uint8)

    for _ in range(num_defect):
        defect_mask_id = np.random.randint(len(defect_mask_path_list))
        mask = cv2.imread(defect_mask_path_list[defect_mask_id], 0)

        # 随机缩放、旋转和扭曲掩膜
        scale = np.random.uniform(scale_range[0], scale_range[1])
        mask_resized = cv2.resize(mask, None, fx=scale, fy=scale, interpolation=cv2.INTER_NEAREST)

        if rotate:
            angle = random.uniform(0, 90)
            h, w = mask_resized.shape[:2]
            center = (w // 2, h // 2)
            M = cv2.getRotationMatrix2D(center, angle, 1.0)
            mask_resized = cv2.warpAffine(mask_resized, M, (w, h), flags=cv2.INTER_LINEAR,
                                          borderMode=cv2.BORDER_CONSTANT, borderValue=0)

        if distort:
            h, w = mask_resized.shape[:2]
            map_x, map_y = np.meshgrid(np.arange(w), np.arange(h))
            map_x = map_x.astype(np.float32)
            map_y = map_y.astype(np.float32)

            distortion_value = random.randint(2, 5)
            map_x += np.random.normal(loc=0, scale=distortion_value, size=map_x.shape).astype(np.float32)
            map_y += np.random.normal(loc=0, scale=distortion_value, size=map_y.shape).astype(np.float32)

            mask_resized = cv2.remap(mask_resized, map_x, map_y, interpolation=cv2.INTER_LINEAR,
                                     borderMode=cv2.BORDER_CONSTANT, borderValue=0)

        is_gen, counter = True, 0
        while is_gen:
            # 获取叠加区域的位置
            mask_h, mask_w = mask_resized.shape[:2]
            start_x = np.random.randint(0, image_w - mask_w + 1)
            start_y = np.random.randint(0, image_h - mask_h + 1)
            white_y, white_x = np.where(mask_resized != 0)

            counter += 1
            # 确定要覆盖的区域
            min_gray, max_gray = gray_range
            rand_color = random.randint(min_gray, max_gray)
            rand_ch = int(np.random.choice(['0', '2']))
            input_img_patch = input_image_cpy[start_y:start_y + mask_h, start_x:start_x + mask_w, :].copy()
            input_img_patch = input_img_patch.astype(np.int16)
            input_img_patch[white_y, white_x, rand_ch] = input_img_patch[white_y, white_x, rand_ch] + rand_color
            new_values = np.clip(input_img_patch, 0, 255)  # 截断到 [0, 255] 范围

            # temp_image[seed_y, seed_x, rand_ch] = input_image_cpy[seed_y, seed_x, rand_ch] - color
            temp_image = input_image_cpy[start_y:start_y + mask_h, start_x:start_x + mask_w]
            mean_val_ori = np.median(temp_image[white_y, white_x, :])
            mean_val_aug = np.median(new_values[white_y, white_x, :])
            if abs(mean_val_ori-mean_val_aug) > 25 and 250 > mean_val_aug > 10:
                is_gen = False
            else:
                is_gen = True
            if counter == 5:
                is_gen = False
        if counter == 5:
            continue
        # 写回图像
        input_image_cpy[start_y:start_y + mask_h, start_x:start_x + mask_w, :] = new_values.astype(np.uint8)
        # 更新缺陷掩膜
        defect_mask[start_y:start_y + mask_h, start_x:start_x + mask_w][white_y, white_x] = mask_resized[white_y, white_x].copy()

    if np.count_nonzero(defect_mask) == 0:
        return None, None, False
    else:
        _, defect_mask = cv2.threshold(defect_mask, 128, 255, cv2.THRESH_BINARY)
        return input_image_cpy, defect_mask, True


def gen_color_mask(rgb_image):
    """
    返回一张纯色区域掩膜mask(0:代表纸张颜色， 255：代表纯色区域)
    """
    bi_image = cv2.bilateralFilter(rgb_image, d=9, sigmaColor=25, sigmaSpace=25)
    channel_b, channel_g, channel_r = cv2.split(bi_image)
    diff_rg = cv2.absdiff(channel_r, channel_g)
    diff_gb = cv2.absdiff(channel_g, channel_b)
    diff_br = cv2.absdiff(channel_b, channel_r)
    max_diff_mask = np.maximum(np.maximum(diff_rg, diff_gb), diff_br)
    max_diff_mask[max_diff_mask > 35] = 255
    return max_diff_mask


def clamp_value(center, offset):
    """
    给定一个中心值和偏移量，计算加或减偏移量后的值，并将其限制在 [0, 255] 范围内。

    :param center: 中心值（int 或 float）
    :param offset: 偏移量（int 或 float），可以是正数也可以是负数
    :return: 限制在 [0, 255] 范围内的结果值（int）
    """
    # 计算新的值
    new_value = center + offset

    # 使用 max 和 min 来限制范围
    clamped_value = max(0, min(255, new_value))

    # 返回整数形式的结果
    return int(clamped_value)


def aug_shift_defect(input_image, shift_x=(-2, 2), shift_y=(-2, 2)):
    """
    说明：对图片进行亚像素平移，仿真配准异常现象
    参数：
    - image_array: 输入图片（灰度图和RGB均可）
    - shift_x: 水平方向的偏移量（正数向右，负数向左）
    - shift_y: 垂直方向的偏移量（正数向下，负数向上）
    返回:
    - 平移后的图片
    """
    # 获取图片的高度和宽度
    h, w = input_image.shape[:2]

    # 定义仿射变换矩阵（平移矩阵）
    min_shift_x, max_shift_x = shift_x
    random_shift_x = np.random.uniform(min_shift_x, max_shift_x)
    min_shift_y, max_shift_y = shift_y
    random_shift_y = np.random.uniform(min_shift_y, max_shift_y)
    translation_matrix = np.float32([[1, 0, random_shift_x], [0, 1, random_shift_y]])

    # 使用cv2.warpAffine进行平移操作
    translated_image = cv2.warpAffine(input_image, translation_matrix, (w, h), borderMode=cv2.BORDER_WRAP)

    return translated_image


def aug_light(src_image, gamma=(0.6, 0.9)):
    """
    说明：相机拍摄的时候整体变浅
    """
    seq = iaa.GammaContrast(gamma)
    aug_light_image = seq(image=src_image)
    return aug_light_image




